# Dataset是自定义数据集需要使用到的基类
from torch.utils.data import Dataset

# DataLoader是下面的代码加载自定义数据集使用的，开始自定义数据集的时候不需要用到
from torch.utils.data import DataLoader

# 数据处理用到，如果训练模型，需要转换成Tensor
from torch import Tensor

# 自定义数据集类
# 数据集中有以下样本：
# | 样本数据     | 样本标签 |
# | :----        | :----    |
# | [0.01, 0.02, 0.03] | [ 0.1 ]    |
# | [0.04, 0.05, 0.06] | [ 0.2 ]    |
# | [0.07, 0.08, 0.09] | [ 0.3 ]    |
# | [0.10, 0.11, 0.12] | [ 0.4 ]    |
# | [0.13, 0.14, 0.15] | [ 0.5 ]    |
class CustomDataset(Dataset):

    # 初始化自定义数据集
    def __init__(self, transform=None, target_transform=None):
        self.transform = transform
        self.target_transform = target_transform

    # __len__方法需要返回数据集中样本的数量
    def __len__(self):
        return 5

    # 数据集在被使用的时候，__getitem__方法会被调用。
    # 参数idx是索引，对应__getitem__需要返回的第idx个样本的数据和对应的标签
    def __getitem__(self, idx):
        dataset = [
            {'data': [0.01, 0.02, 0.03], 'label': [ 0.1 ]},
            {'data': [0.04, 0.05, 0.06], 'label': [ 0.2 ]},
            {'data': [0.07, 0.08, 0.09], 'label': [ 0.3 ]},
            {'data': [0.10, 0.11, 0.12], 'label': [ 0.4 ]},
            {'data': [0.13, 0.14, 0.15], 'label': [ 0.5 ]}
        ]
        data = dataset[idx]['data']
        label = dataset[idx]['label']
        if self.transform:
            data = self.transform(data)
        if self.target_transform:
            label = self.target_transform(label)
        #return Tensor(data), Tensor(label)
        return data, label

# batch_size影响下面每一次for语句从test_dataloader取出的data和label中的tensor数据大小
batch_size = 10

# 创建数据集对象和DataLoader对象
test_data = CustomDataset(Tensor, Tensor)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

# 遍历和打印数据
for data, label in test_dataloader:
    print('data', data)
    print('label', label)

# 输出：
# data tensor([[0.0100, 0.0200, 0.0300],
#         [0.0400, 0.0500, 0.0600],
#         [0.0700, 0.0800, 0.0900],
#         [0.1000, 0.1100, 0.1200],
#         [0.1300, 0.1400, 0.1500]])
# label tensor([[0.1000],
#         [0.2000],
#         [0.3000],
#         [0.4000],
#         [0.5000]])
